"""Simple MNIST image classifier example with LightningModule and LightningDataModule.

To run: python minmst_example.py
"""

import os
import torch
import torch_npu

from examples.dataset_moudle.mnist_datamodule import MNISTDataModule
from examples.module.mnist_module import ImageClassifier
from pytorch_lightning.utilities.cli import LightningCLI
from pytorch_lightning.callbacks import ModelCheckpoint

from lightning_npu.accelerators.npu import NPUAccelerator
from lightning_npu.strategies.npu_parallel import NPUParallelStrategy


def _get_last_ckpt(ckpt_dir):
    ckpt_files = [ckpt_file for ckpt_file in os.listdir(ckpt_dir)
                  if ckpt_file.endswith('.ckpt')]
    if not ckpt_files:
        print("No ckpt file found.")
        return None
    ckpt_files = sorted(ckpt_files,key=lambda x: os.path.getatime(os.path.join(ckpt_dir, x)))
    return os.path.join(ckpt_dir, ckpt_files[-1])


def cli_main():
    # The LightningCLI removes all the boilerplate associated with arguments parsing. This is purely optional.
    checkpoint_callback = ModelCheckpoint(dirpath="./ckpt/", 
                                          save_top_k=2, 
                                          monitor="val_loss",
                                          mode="min",
                                          filename="sample-mnist--{epoch:02d}-{val_loss:.2f}"
                                          )
    cli = LightningCLI(
        ImageClassifier,
        MNISTDataModule,
        trainer_defaults={
            "accelerator": NPUAccelerator(),
            "devices": 8,
            "max_epochs": 5,
            "strategy": NPUParallelStrategy(),
            "callbacks": [checkpoint_callback],
        },
        seed_everything_default=42,
        save_config_overwrite=True,
        run=False
    )
    cli.trainer.fit(cli.model, datamodule=cli.datamodule)
    cli.trainer.save_checkpoint("myckpt.ckpt")
    cli.trainer.test(ckpt_path="best", datamodule=cli.datamodule)


    checkpoint_callback = ModelCheckpoint(dirpath="./ckpt/", 
                                          save_top_k=1, 
                                          monitor="val_loss",
                                          mode="min",
                                          filename="sample-mnist--{epoch:02d}-{epoch}",
                                          save_weights_only=True,
                                          )
    cli = LightningCLI(
        ImageClassifier,
        MNISTDataModule,
        trainer_defaults={
            "accelerator": NPUAccelerator(),
            "devices": 8,
            "max_epochs": 8,
            "strategy": NPUParallelStrategy(),
            "callbacks": [checkpoint_callback],
        },
        seed_everything_default=42,
        save_config_overwrite=True,
        run=False
    )
    checkpoint = _get_last_ckpt("./ckpt/")
    cli.trainer.fit(cli.model, datamodule=cli.datamodule, ckpt_path=checkpoint)
    checkpoint_callback.best_model_path
    cli.trainer.test(ckpt_path="best", datamodule=cli.datamodule)


if __name__ == "__main__":
    cli_main()